import matplotlib.pyplot as plt
import numpy as np
import os


def figure_save_curve(path, x, y1, y2, title, legend_name, epochs):
    """
    :param
    path:the position to store figures
    x:is the epochs of training
    y1:is the testing result on train dataset
    y2:is the testing result on validation dataset
    legend_name:to describe the legend of curves such as ["train_loss","val_loss"]
    """
    plt.figure()
    plt.plot(x, y1)
    plt.plot(x, y2)
    plt.title(title)
    plt.legend(legend_name)
    plt.yticks(np.arange(0, 1.05, 0.05))
    plt.savefig(f"{path}/{title}_{epochs}epochs")


def figure_all_metric(path, name, train_loss, val_loss, train_acc, train_precision, train_recall, train_f1,
                      val_acc, val_precision, val_recall, val_f1, epochs):
    path = os.path.join(path, name)
    if not os.path.exists(path):
        os.mkdir(path)

    num_epochs = len(train_loss)
    x = list(range(1, num_epochs + 1))
    # 画loss图
    figure_save_curve(path, x, train_loss, val_loss, "loss", ["train_loss", "val_loss"], epochs)
    # 画acc图
    figure_save_curve(path, x, train_acc, val_acc, "accuracy", ["train_acc", "val_acc"], epochs)
    # plt.ylim(0, 40)
    # plt.yticks(np.arange(70, 101, 2))
    # 画precision图
    figure_save_curve(path, x, train_precision, val_precision, f"{name}-precision",
                      ["train_precision", "val_precision"],
                      epochs)
    # 画recall图
    figure_save_curve(path, x, train_recall, val_recall, f"{name}_recall", ["train_recall", "val_recall"], epochs)
    # 画f1score图
    figure_save_curve(path, x, train_f1, val_f1, f"{name}-F1_score", ["train_f1", "val_f1"], epochs)


def show_result(path, train_loss, val_loss, train_acc_dict, train_precision_dict, train_recall_dict, train_f1_dict,
                val_acc_dict, val_precision_dict, val_recall_dict, val_f1_dict, epochs):
    if not os.path.exists(path):
        os.mkdir(path)
    data_name = ["total", "hole", "crack", "normal"]
    for name in data_name:
        with open(f"{path}/train_info.txt", "a+") as f:
            f.writelines([
                f"{name}_train_loss:{list(train_loss)}" + '\n',
                f"{name}_train_acc:{list(train_acc_dict[name])}" + '\n',
                f"{name}_train_precision:{list(train_precision_dict[name])}" + '\n',
                f"{name}_train_recall:{list(train_recall_dict[name])}" + '\n',
                f"{name}_train_f1{list(train_f1_dict[name])}" + '\n'
            ])
        with open(f"{path}/val_info.txt", "a+") as f:
            f.writelines([
                f"{name}_val_loss:{list(val_loss)}" + '\n',
                f"{name}_val_acc:{list(val_acc_dict[name])}" + '\n',
                f"{name}_val_precision:{list(val_precision_dict[name])}" + '\n',
                f"{name}_val_recall:{list(val_recall_dict[name])}" + '\n',
                f"{name}_val_f1:{list(val_f1_dict[name])}" + '\n'
            ])

        figure_all_metric(path, name, train_loss, val_loss, train_acc_dict[name], train_precision_dict[name],
                          train_recall_dict[name],
                          train_f1_dict[name], val_acc_dict[name], val_precision_dict[name], val_recall_dict[name],
                          val_f1_dict[name], epochs)
